import cv2
import numpy as np
from onnx_infer.cartoon.cartoon import Cartoon
from onnx_infer.detect.OnnxFaceDetector import UltraLightFaceDetection
from onnx_infer.config import GetConfig
from onnx_infer.utils.utils import mask_merge

class Ensembel:
    def __init__(self,mask_image_path,config) -> None:
        self.args = config
        
        self.mask_image = self._get_mask_image(mask_image_path)

        self.detect = UltraLightFaceDetection(self.args.detect_onnx)
        self.cartoon_256 = Cartoon(self.args.cartoon_onnx_256)
        self.cartoon_512 = Cartoon(self.args.cartoon_onnx_512)
    @staticmethod
    def _get_mask_image(image_path):
        mask_image = cv2.imread(image_path)
        return cv2.cvtColor(mask_image,cv2.COLOR_BGR2GRAY)
    
    def run2(self,image_path):
        image = cv2.imread(image_path)
        bboxes,_ = self.detect.inference(image)
        image_bg = self.cartoon_512.run(image,512)
        h,w,_ = image.shape
        extension_factor = 0.1
        
        for bbox in bboxes.astype(int):
            left = bbox[0]
            up = bbox[1]
            right = bbox[2]
            down = bbox[3]
            
            width = right-left
            height = down-up
            
            left = max(0,left-int(width*extension_factor))
            up   = max(0,up-int(height*extension_factor))
            right= min(w,int(right+width*extension_factor))
            down = min(h,int(down+height*extension_factor))
            # print(up,down,left,right)
            # cv2.imwrite("result/result2.jpg",image[up:down,left:right,...])
            sub_image = self.cartoon_256.run(image[up:down,left:right,...],256)
            
            sub_image_cartoon = image_bg[up:down,left:right,...]
            merged_image = mask_merge(self.mask_image,sub_image,sub_image_cartoon)
            image_bg[up:down,left:right,...] = merged_image
            # image_bg=cv2.rectangle(image_bg,(left,up),(right,down),(255,0,0),2)
        
        return image_bg
    
if __name__ == "__main__":
    file_path = 'asset/37.jpg'
    mask_path = 'utils/mask_images/mask1.jpg'
    model = Ensembel(mask_image_path=mask_path)
    image = model.run2(file_path)
    
    cv2.imwrite("result/result.jpg",image)
    
    import os
    import numpy as np
    base_dir = 'asset'
    for file in os.listdir(base_dir):
        file_path = os.path.join(base_dir,file)
        mask_path = 'utils/mask_images/mask1.jpg'
        model = Ensembel(mask_image_path=mask_path)
        output_img = model.run2(file_path)
        input_img = cv2.imread(file_path)
        res = np.concatenate([output_img,input_img])
        cv2.imwrite("result/"+file,res)
    